package com.dal.guava.aspect;

import com.google.common.util.concurrent.RateLimiter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;

import java.time.Instant;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.TimeUnit;

/**
 * 高性能分布式限流组件
 * 结合本地令牌桶和Redis实现精准限流
 */
@Component
public class DistributedRateLimiter {
    private static final Logger log = LoggerFactory.getLogger(DistributedRateLimiter.class);


    // 本地令牌桶缓存（用于快速判断）
    private final ConcurrentMap<String, RateLimiter> localLimiters = new ConcurrentHashMap<>();
    
    // Redis key前缀
    private static final String REDIS_KEY_PREFIX = "rate_limit:";
    
    // 滑动窗口限流Lua脚本
    private static final RedisScript<Long> SLIDING_WINDOW_SCRIPT;

    // 更新后的Lua脚本
    // 更新Lua脚本，更宽松的参数校验
    static {
        StringBuilder script = new StringBuilder();
        script.append("local key = KEYS[1]");
        script.append("\nlocal timestamp = tonumber(ARGV[1] or '')");  // 处理空字符串
        script.append("\nlocal window = tonumber(ARGV[2] or '')");
        script.append("\nlocal limit = tonumber(ARGV[3] or '')");

        // 允许limit为小数，但实际使用时取整数
        script.append("\nlimit = math.floor(limit or 0)");

        // 参数校验（允许0，但不允许nil）
        script.append("\nif timestamp == nil or window == nil or limit == nil then");
        script.append("\n  return -1");  // 参数为nil时返回-1
        script.append("\nend");

        // 窗口和限流值必须为正数
        script.append("\nif window <= 0 or limit <= 0 then");
        script.append("\n  return -2");  // 参数无效时返回-2
        script.append("\nend");

        // 核心限流逻辑
        script.append("\nlocal current = redis.call('zcount', key, timestamp - window, timestamp)");
        script.append("\nif current + 1 > limit then");
        script.append("\n  return 0");  // 限流
        script.append("\nelse");
        script.append("\n  redis.call('zadd', key, timestamp, timestamp)");
        script.append("\n  redis.call('zremrangebyscore', key, 0, timestamp - window)");
        script.append("\n  redis.call('expire', key, math.ceil(window / 1000) + 1)");  // 确保过期时间为整数
        script.append("\n  return 1");  // 允许通过
        script.append("\nend");

        SLIDING_WINDOW_SCRIPT = new DefaultRedisScript<>(script.toString(), Long.class);
    }

    @Autowired
    private RedisTemplate<String, Object> redisTemplate;
    
    /**
     * 尝试获取令牌（先本地后Redis）
     * @param key 限流标识
     * @param permitsPerSecond 每秒允许的请求数
     * @param timeout 超时时间
     * @param unit 时间单位
     * @return 是否获取成功
     */
    public boolean tryAcquire(String key, double permitsPerSecond, long timeout, TimeUnit unit) {
        // 1. 先尝试从本地令牌桶获取（快速判断）
        RateLimiter localLimiter = localLimiters.computeIfAbsent(key, 
            k -> RateLimiter.create(permitsPerSecond));
        
        if (localLimiter.tryAcquire(timeout, unit)) {
            // 2. 本地获取成功后，再到Redis进行精确校验
            return checkInRedis(key, permitsPerSecond);
        }
        
        return false;
    }
    
    /**
     * 在Redis中进行精确限流校验
     * @param key 限流标识
     * @param permitsPerSecond 每秒允许的请求数
     * @return 是否允许请求
     */
    private boolean checkInRedis(String key, double permitsPerSecond) {
        String redisKey = REDIS_KEY_PREFIX + key;
        long currentTime = Instant.now().toEpochMilli();
        long window = 1000; // 1秒窗口
        long limit = (long) permitsPerSecond; // 强制转换为长整数，避免小数

        // 日志打印参数，确认格式正确
//        log.debug("参数检查: timestamp={}({}), window={}({}), limit={}({})",
//                currentTime, currentTime.getClass().getSimpleName(),
//                window, window.getClass().getSimpleName(),
//                limit, limit.getClass().getSimpleName());

        // 执行脚本时使用长整数参数
        List<String> keys = Collections.singletonList(redisKey);
        List<String> args = Arrays.asList(
                String.valueOf(currentTime),  // 纯数字字符串，如"1753613954955"
                String.valueOf(window),       // 纯数字字符串，如"1000"
                String.valueOf(limit)         // 纯数字字符串，如"1"（避免小数）
        );

        Long result = redisTemplate.execute(SLIDING_WINDOW_SCRIPT, keys, args.toArray());

        // 处理返回结果
        if (result == null) {
            log.error("Redis脚本执行失败，返回null");
            return false;
        } else if (result == -1) {
            log.error("Redis脚本参数错误，检查参数格式");
            return false;
        }
        return result == 1;
    }
    
    /**
     * 重置限流状态
     */
    public void reset(String key) {
        localLimiters.remove(key);
        redisTemplate.delete(REDIS_KEY_PREFIX + key);
    }
}